#include "../Base.h"


typedef NTSTATUS(NTAPI* KdEnumerateDebuggingDevices)(PVOID un1, PVOID un2, PVOID un3);

KdEnumerateDebuggingDevices oldKdEnumerateDebuggingDevices = NULL;

NTSTATUS newKdEnumerateDebuggingDevices(PVOID info, PVOID un2, PVOID un3) {

	if (MmIsAddressValid(info)) {
		PCommPackage package = (PCommPackage)info;
		if (package->id == 0x12345678) {
			package->retStatus = gCommCallback(package);
		}
		else {

			if (oldKdEnumerateDebuggingDevices) {
				return oldKdEnumerateDebuggingDevices(info, un2, un3);
			}
		}
	}
	return STATUS_SUCCESS;
}

NTSTATUS RegisterCommWin10(CommCallbackProc callback) {

	PUCHAR funcAddress = (PUCHAR)searchCode("ntoskrnl.exe", "PAGE", "488B05****75*488B05****", 0);
	if (!funcAddress) {
		LOG("search error\r\n");
		return STATUS_UNSUCCESSFUL;
	}
	ULONG64 offset = *(PLONG)(funcAddress  + 3);
	PULONG64 repFunc = (PULONG64)((funcAddress+ 7) + offset);
	oldKdEnumerateDebuggingDevices = (KdEnumerateDebuggingDevices)repFunc[0];
	repFunc[0] = (ULONG64)newKdEnumerateDebuggingDevices;
	gWin10CommCallBack = repFunc;
	gCommCallback = callback;
	return STATUS_SUCCESS;
}



VOID UnRegisterCommWin10() {

	if (oldKdEnumerateDebuggingDevices) {
		gWin10CommCallBack[0] = (ULONG64)oldKdEnumerateDebuggingDevices;
	}
}
